from pylab import *
from scipy.special import dawsn
import scipy.ndimage

def myf(x):
    return 4.0/sqrt(pi) * (-x -dawsn(x) +2*x**2*dawsn(x))

def myf_maybe(x):
    return (-2.254*x+x*x*x)*exp(-x*x)

ion()

Nw = 11  ## Window size
Nf = 512 ## number of different angles to calculate
# Np = 120 ## size of syntetic image to generate

#sigma1 = 1.5
sigma1 = 3.0
#sigma1 = 4.5
sigma2 = 1.5

yyo,xxo = mgrid[0.:Nw,0.:Nw] / floor(Nw/2) - 1.0
theta = arctan2(yyo,xxo)
yy,xx = zeros((2,Nw,Nw))

bank_c = 1j*zeros((Nf,Nw,Nw))
for n in range(Nf):
    phi = 2*pi*n/(Nf)
    print 'phi:', phi
    R = array([[cos(phi), -sin(phi)],[sin(phi),cos(phi)]])
    ## Generate rotated vectors
    for j in range(Nw):
        for k in range(Nw):
            xx[j,k],yy[j,k] = dot(R, [xxo[j,k],yyo[j,k]])

    bank_c[n] = (exp(-(sigma2*yy)**2) * exp(-(sigma1*xx)**2) * (4 * (sigma1*xx)**2 - 2)) \
                + 1j * (exp(-(sigma2*yy)**2) * myf(sigma1*xx))


# ## Create circular test image
# img = zeros((Np,Np))
# for j in range(Np):
#     for k in range(Np):
#         d = sqrt(((j-Np/2.)/Np)**2 + ((k-Np/2.)/Np)**2)
#         if d < .4 and d>.2:
#             img[j,k] = 1

## The products of filters with different directions.
h_t = 1j*zeros(Nf)
for n in range(Nf):
    h_t[n] = np.sum(bank_c[n]*conj(bank_c[0]))

## Find maximum filter frequency we care about.
H_t = np.abs(fft(real(h_t)))

err_lim = 1e-3
Nfe = nonzero(H_t< (H_t[0] * err_lim) )[0][0]
print 'Nfe:',Nfe
## Forcing this
#Nfe = 24

## Calculate the sigma_i and a_i functions
sigma_i = sqrt(H_t)
a_i = 1j * zeros((Nfe*2-1,Nw,Nw))
deltatheta = 2*pi/(Nf)

n=0
nu_i = n
for m in range(Nf):
    theta = m*deltatheta
    a_i[n] += bank_c[m]
a_i[n] = a_i[n] * deltatheta / sigma_i[n]/(2*pi)
#a_i[n] = a_i[n] * deltatheta / (2*pi)
a_i/=2

for n in range(1,Nfe):
    nu_i = n
    for m in range(Nf):
        theta = m*deltatheta
        a_i[n] += bank_c[m]*exp(1j * nu_i * theta)
        a_i[-n] += bank_c[m]*exp(1j * -nu_i * theta)
    a_i[n] = a_i[n] * deltatheta / sigma_i[n]/(2*pi)
    a_i[-n] = a_i[-n] * deltatheta / sigma_i[n]/(2*pi)
    # a_i[n] = a_i[n] * deltatheta / (2*pi)
    # a_i[-n] = a_i[-n] * deltatheta / (2*pi)

## The reconstructed filters
hbank_c = 1j*zeros((Nf,Nw,Nw))
for n in range(Nf):
    theta = n*deltatheta
    for m in range(Nfe):
        hbank_c[n] += sigma_i[m] * a_i[m] * exp(1j * -m * theta)
        hbank_c[n] += (-1)**m * sigma_i[m] * conj(a_i[m]) * exp(1j * m * theta)
        #hbank_c[n] += a_i[m] * exp(1j * -m * theta)
        #hbank_c[n] += (-1)**m * conj(a_i[m]) * exp(1j * m * theta)



figure(1)
suptitle('Original filters')
for n in range(32):
    subplot(8,8,n+1)
    #pp = imshow(bank_e[n], vmin=-2, vmax=2, extent=(-1,1,-1,1))#, cmap=cm.bone)
    pp = imshow(real(bank_c[n]), vmin=-2, vmax=2, extent=(-1,1,-1,1))#, cmap=cm.bone)
    pp.get_axes().get_xaxis().set_ticks([])
    pp.get_axes().get_yaxis().set_ticks([])

    subplot(8,8,n+32+1)
    #pp = imshow(bank_o[n], vmin=-2, vmax=2, extent=(-1,1,-1,1))#, cmap=cm.bone)
    pp = imshow(imag(bank_c[n]), vmin=-2, vmax=2, extent=(-1,1,-1,1))#, cmap=cm.bone)
    pp.get_axes().get_xaxis().set_ticks([])
    pp.get_axes().get_yaxis().set_ticks([])

figure(2)
subplot(2,1,1)
plot(real(h_t), '-+')
plot(imag(h_t), '-')
subplot(2,1,2)
semilogy(np.abs(fft(real(h_t))), '-+')

figure(3)
semilogy(sigma_i/sigma_i[0], '-o')
grid()

figure(4)
ll = 0.01
suptitle('A_i functions')
for n in range(10):
    subplot(4,5,n+1)    
    pp = imshow(real(a_i[n]), vmin=-ll, vmax=ll, extent=(-1,1,-1,1))
    pp.get_axes().get_xaxis().set_ticks([])
    pp.get_axes().get_yaxis().set_ticks([])

    subplot(4,5,n+10+1)
    pp = imshow(imag(a_i[n]), vmin=-ll, vmax=ll, extent=(-1,1,-1,1))
    pp.get_axes().get_xaxis().set_ticks([])
    pp.get_axes().get_yaxis().set_ticks([])

figure(5)
suptitle('Approximated filters')
for n in range(32):
    subplot(8,8,n+1)
    pp = imshow(real(hbank_c[n]), vmin=-2, vmax=2, extent=(-1,1,-1,1))
    pp.get_axes().get_xaxis().set_ticks([])
    pp.get_axes().get_yaxis().set_ticks([])

    subplot(8,8,n+32+1)
    pp = imshow(imag(hbank_c[n]), vmin=-2, vmax=2, extent=(-1,1,-1,1))
    pp.get_axes().get_xaxis().set_ticks([])
    pp.get_axes().get_yaxis().set_ticks([])

figure(6)
suptitle('Difference')
for n in range(32):
    subplot(8,8,n+1)
    pp = imshow(real(bank_c[n]-hbank_c[n]), vmin=-1e-1, vmax=1e-1, extent=(-1,1,-1,1))
    pp.get_axes().get_xaxis().set_ticks([])
    pp.get_axes().get_yaxis().set_ticks([])

    subplot(8,8,n+32+1)
    pp = imshow(imag(bank_c[n]-hbank_c[n]), vmin=-1e-1, vmax=1e-1, extent=(-1,1,-1,1))
    pp.get_axes().get_xaxis().set_ticks([])
    pp.get_axes().get_yaxis().set_ticks([])

